"""
author: liuxu
date: 2025/5/30
description: 复测hipprof trace相关的bug
"""
import os
import re
import allure
import pytest
from collections import Counter, defaultdict

from common.logger_ctrl import mylogger
from common.ssh_command import exec_cmd
from common.get_hardware_info import get_system_info
from base_public.testbase import BasePublic
from setting import CPP_DIR, PROJECT_ROOT


@pytest.mark.low
@pytest.mark.performance
@allure.feature('hip trace相关bug')
class TestTraceBugCase(BasePublic):
    @pytest.mark.timeout(600)
    def test_bug_78850(self, test_env):
        """
        Bug链接: http://hpczentao.sugon.com/bug-view-78850.html
        标题: hipprof记录API时间戳流程导致主机端产生10~40us空泡
        问题描述: trace的可视化图里查看kernel调用空泡较大   --> 此bug先后转为24.10，25.04的需求

        需求链接：http://hpczentao.sugon.com/story-view-8531-1-project-583.html
        验收标准：hygonCPU与intelCPU环境下，hipprof单次时间测量（开始+结束）对主机端应用损耗小于2us
        :return:
        """
        # 前置条件：获取cpu型号
        expect_bubble_time = 2 * (10 ** 3)  # hygon cpu环境空泡小于2us
        system_info = get_system_info(categories=["cpu"])
        cpu_model = system_info["cpu"]["model"].lower()
        mylogger.debug(f"cpu_model: {cpu_model}")
        if "intel" in cpu_model:
            expect_bubble_time = 1000  # intel cpu环境空泡小于1us

        # 编译用例
        src_case_dir = os.path.join(CPP_DIR, "hipSimple/hipDispatchLatency")
        code_path = self.compile_source_code(src_file_name=os.path.join(src_case_dir, "test_kernel.cpp"),
                                             des_exe_name="test_kernel.code", compile_paras="--genco")
        ret0, _ = exec_cmd(f"ln -sf {code_path} ./test_kernel.code")  # 解决工作目录与code所在目录不一致导致用例执行失败问题
        exe_path = self.compile_source_code(src_file_name=os.path.join(src_case_dir, "hipDispatchEnqueueRateMT.cpp"),
                                            des_exe_name="hipDispatchEnqueueRateMT.out")

        # 执行测试
        # shell_jq_path = os.path.join(PROJECT_ROOT, "shell_script/analyze_trace.sh")
        py_analyze_trace_json = os.path.join(PROJECT_ROOT, "common/analyze_trace.py")
        for thread_num in (5, 10, 20):
            cmd_timeout = thread_num * 30  # 设置命令执行的超时时间
            ## 将--segment-size设置较大，以保证json不分割
            # ret, output = exec_cmd(f"hipprof --hip-trace --segment-size {500 * 1000} {exe_path} {thread_num} 1-hipModuleLaunchKernel", cwd=test_env)
            ret, output = exec_cmd(f"hipprof --hip-trace --trace-args {exe_path} {thread_num} 1-hipModuleLaunchKernel", cwd=test_env, timeout=cmd_timeout)
            assert ret is True
            # 获取json文件路径，使用shell脚本结合jq命令分析
            json_names = " ".join(
                [json_info["name"] for json_info in self.get_all_prof_files_info(output)["json_list"]])  # json个数>=1
            ret_, output_ = exec_cmd(f"python3 {py_analyze_trace_json} {json_names}", cwd=test_env, timeout=cmd_timeout)
            assert ret_ is True

            # 校验每个线程执行的平均空泡时间
            bubble_data_dict = defaultdict(list)  # 存放处理后的空泡数据
            for data_line in output_.split("\n"):
                # 去除打印的标题行,以及不匹配的事件行
                if data_line.startswith("TID") or data_line.startswith("==="):
                    continue
                tid,total_count,first,max_,second_max,min_,avg,avg_excl_first = [float(i) for i in data_line.split()]
                bubble_data_dict[tid] = [total_count, avg]

            # 校验测试结果
            assert len(bubble_data_dict) == thread_num
            for tid, (_, avg_bubble_time) in bubble_data_dict.items():
                assert avg_bubble_time < expect_bubble_time, f"The result of thread {tid} does not meet expectations"

    def test_bug_71772(self, test_env, dcu_info):
        """
        Bug链接: http://hpczentao.sugon.com/bug-view-71772.html
        标题: 【HPC-客户应用】hipprof在CUDA_VISIBLE_DEVICES场景下不能正确识别真正的deviceid
        问题原因: hipprof获取的多卡用例显示id是针对软件层面的编号，所以每个进程显示的卡号都是0
        修改方案：用device的uuid来区分卡，显示卡的实际序号与hy-smi保持一致
        :return:
        """
        if dcu_info.get("count", 0) < 4:
            pytest.skip(
                "The test environment does not meet the test conditions(The script 'run.sh' needs to start four gemms running on different cards),Skip the test case!")
        # 编译用例
        src_case_dir = os.path.join(CPP_DIR, "bug_case/bug-71772")
        code_path = self.compile_source_code(src_file_name=os.path.join(src_case_dir, "demo_gemm_nvcc.cu"),
                                             des_exe_name="v100_demo_gemm",
                                             compiler="nvcc",
                                             compile_paras="-I$ROCM_PATH/cuda/include -L$ROCM_PATH/cuda/lib64 -lcublas -lcurand -lpthread")
        exec_cmd(f"ln -sf {code_path} ./v100_demo_gemm")  # 解决工作目录与脚本目录不一致导致用例执行失败问题
        run_sh_path = os.path.join(src_case_dir, "run.sh")

        # 执行trace测试,demo脚本内部通过CUDA_VISIBLE_DEVICES控制程序运行在不同卡上
        ret, output = exec_cmd(f"hipprof bash {run_sh_path}", cwd=test_env, timeout=60)
        assert ret is True
        json_path = self.get_all_prof_files_info(output)["json_list"][0]["name"]

        device_id_list = []
        device_num_pattern = re.compile(r'\[3\] Compute on Device (\d+)')
        for trace_item_d in self.parse_prof_json(json_path):
            if trace_item_d.get("ph") == "X":
                break
            category_name = trace_item_d.get("args", {}).get("name", "")

            if r"[3] Compute on Device" in category_name:
            #     device_id_list.append(int(category_name.replace("[3] Compute on Device", "")))
                match = device_num_pattern.search(category_name)
                if match:
                    # 从捕获组 1 中提取数字（转为 int 更实用）
                    device_num = int(match.group(1))
                    device_id_list.append(device_num)
                else:
                    pytest.fail(f"The regular expression failed to obtain the device id in this category name({category_name})")

        assert len(device_id_list) == len(set(device_id_list)) > 1

    def test_bug_93269(self, test_env):
        """
        Bug链接: http://hpczentao.sugon.com/bug-view-93269.html
        标题: 【DTK-25.04.1-alpha-0421】hipprof命令使用报错问题：字符串转义或使用不当造成的问题
        问题描述：hipprof问题，处理字符转义bug
        :return:
        """
        exe_path = os.path.join(CPP_DIR, "bug_case/bug-93269/bug93269_test.sh")
        special_characters = '")(*&^|]}中文"  "para2"'
        ret, output = exec_cmd(f"hipprof --libc-trace bash {exe_path} {special_characters}")
        assert ret is True

        # 检测test.sh脚本打印的命令行参数是否发生转义
        for line in output.split("\n"):
            if "input paras:" in line:
                assert special_characters in line
        # 检测是否正常生成trace数据
        libc_csv_path = self.get_all_prof_files_info(output, test_env)["libc_api_csv"]["abspath"]
        libc_trace_d = self.get_trace_info_from_csv(libc_csv_path)
        assert len(libc_trace_d) > 0

    def test_bug_93852(self, test_env):
        """
        Bug链接: http://hpczentao.sugon.com/bug-view-93852.html
        标题: 【dtk-25.04.1-beta-0426-centos7】--libc-trace 功能执行某些命令存在崩溃情况，-o指令指定同一个名称连续执行两次，存在数据库错误
        问题描述：hipprof问题，--libc-trace问题
        :return:
        """
        # 先清理当前环境文件
        exec_cmd(f"rm -f abc.db", cwd=test_env)
        exec_cmd(f"ls", cwd=test_env)
        # 步骤1
        ret, output = exec_cmd(f"hipprof --libc-trace -o abc ps -aux",  cwd=test_env)
        assert ret is True
        # 验证生成的文件名均为abc
        assert self.exist_file("abc.db")
        assert self.exist_file("abc.libctrace.csv")
        assert self.exist_file("abc.json")

        # 步骤2:校验输出文件已存在时工具不崩溃，并给出提示
        ret2, output2 = exec_cmd(f"hipprof --libc-trace -o abc hy-smi", cwd=test_env)
        assert ret2 is False
        assert "File abc.db already exists,please delete it first" in output2

    # def test_bug_graph_99154(self, test_env):
    #     """
    #     Bug链接: http://hpczentao.sugon.com/bug-view-99154.html
    #     标题:  dtk25.04.1+trorch2.5 torchprof json文件kernel有缺失
    #     问题描述：graph与hipprof对接问题：
    #         galaxy和hipprof中对kernel name的使用或要求存在出入。prof需要整个程序声明周期的char*指针, galaxy传入的是一个非全局string.c_str()
    #     :return:
    #     """
    #     # （1）编译用例
    #     src_case_dir = os.path.join(CPP_DIR, "bug_case/bug-99154")
    #     exe_path = self.compile_source_code(src_file_name=os.path.join(src_case_dir, "hipGraphTest.cpp",),
    #                                         compile_paras="--gpu-max-threads-per-block=1024")
    #
    #     # （2）执行测试
    #     ret, output = exec_cmd(f"hipprof --hip-trace {exe_path}",  cwd=test_env)
    #     assert ret is True
    #
    #     # （3）校验结果
    #     # 验证kernel的顺序是否准确(reduce -> alpha -> beta -> gamma -> delta -> epsilon)
    #     expect_kernel_seq = ["reduce(float*, double*)",
    #                          "kernel_alpha(float*, int)",
    #                          "kernel_beta(float*, int)",
    #                          "kernel_gamma(float*, int)",
    #                          "kernel_delta(float*, int)",
    #                          "kernel_epsilon(float*, int)",
    #                          ]
    #
    #     pid_ = self.get_prof_pid(output)
    #     json_path = f"result_{pid_}.json"
    #     kernel_csv_path = f"result_{pid_}.kernel.csv"
    #     target_kernel_name_l = list(self.get_trace_info_from_csv(kernel_csv_path).keys())
    #     kernel_timings = []
    #     for trace_event in self.parse_prof_json(json_path):
    #         if trace_event.get("name") in expect_kernel_seq:
    #             kernel_name = trace_event["name"]
    #             start_time = int(trace_event["args"].get("BeginNs"))  # 开始时间戳
    #             end_time = int(trace_event["args"].get("EndNs"))  # 结束时间戳
    #             kernel_timings.append({
    #                 "name": kernel_name,
    #                 "start_time": start_time,
    #                 "end_time": end_time
    #             })
    #     # 按hipprof获取的kernel按开始时间排序（确保顺序与执行顺序一致）
    #     kernel_timings.sort(key=lambda x: x["end_time"])
    #
    #     mylogger.debug(f"trace kernel info:\n{self.format_str_print(kernel_timings)}")
    #     prof_kernel_seq = [kernel_timing["name"] for kernel_timing in kernel_timings]
    #
    #     assert expect_kernel_seq == prof_kernel_seq  #TO DO在旧版本还没复现问题

    @pytest.mark.medium
    def test_bug_graph_99154(self, test_env):
        """
        Bug链接: http://hpczentao.sugon.com/bug-view-99154.html
        标题:  dtk25.04.1+trorch2.5 torchprof json文件kernel有缺失
        问题描述：graph与hipprof对接问题：
            galaxy和hipprof中对kernel name的使用或要求存在出入。prof需要整个程序声明周期的char*指针, galaxy传入的是一个非全局string.c_str()
        :return:
        """
        # （1）编译用例
        src_case_dir = os.path.join(CPP_DIR, "bug_case/bug-99154")
        exe_path = self.compile_source_code(src_file_name=os.path.join(src_case_dir, "hipGraphTest_largeOfKernels.cpp",),
                                            compile_paras="--gpu-max-threads-per-block=1024")

        # （2）执行测试
        ret, output = exec_cmd(f"hipprof --hip-trace {exe_path}",  cwd=test_env)
        assert ret is True

        # （3）校验结果
        # 验证kernel的顺序是否准确(reduce -> alpha -> beta -> gamma -> delta -> epsilon)
        expect_kernel_seq = [f"kernel_name{i}(float*, float*, float*, unsigned long)" for i in range(1, 81)]

        pid_ = self.get_prof_pid(output)
        json_path = f"result_{pid_}.json"
        # kernel_csv_path = f"result_{pid_}.kernel.csv"
        # target_kernel_name_l = list(self.get_trace_info_from_csv(kernel_csv_path).keys())
        kernel_timings = []
        for trace_event in self.parse_prof_json(json_path):
            if trace_event.get("name") in expect_kernel_seq:
                kernel_name = trace_event["name"]
                kernel_args = trace_event.get("args")
                assert kernel_args is not None, f"The parameter of kernel is empty, trace_event:\n{trace_event}"
                start_time = int(kernel_args.get("BeginNs") or kernel_args.get("beginNs"))  # 开始时间戳
                end_time = int(kernel_args.get("EndNs") or kernel_args.get("endNs"))  # 结束时间戳
                kernel_timings.append({
                    "name": kernel_name,
                    "start_time": start_time,
                    "end_time": end_time
                })
        # 按hipprof获取的kernel按开始时间排序（确保顺序与执行顺序一致）
        kernel_timings.sort(key=lambda x: x["start_time"])

        mylogger.debug(f"trace kernel info:\n{self.format_str_print(kernel_timings)}")
        prof_kernel_seq = [kernel_timing["name"] for kernel_timing in kernel_timings]

        assert expect_kernel_seq == prof_kernel_seq


@pytest.mark.medium
@allure.feature('omp trace相关bug')
class TestOmpTraceBugCase(BasePublic):
    def test_bug_86580(self, test_env, dcu_info):
        """
        Bug链接: http://hpczentao.sugon.com/bug-view-86580.html
        标题: hipprof中核函数的执行指向了核函数外的内存拷贝
        问题描述: omp trace问题（数据流指向错误是两个事件流的id一致导致的：相同id的两个事件流中：第一组id为14+1拼接pid(452)，第二组id为10+5拼接pid(452)）
        :return:
        """
        # 编译用例
        gfx_arch = dcu_info["gfx_arch"]
        omp_case_path = os.path.join(CPP_DIR, "bug_case/bug-86580/main.f90")
        exe_path = self.compile_source_code(src_file_name=omp_case_path, des_exe_name="main_f90",
                                            compiler="flang",
                                            compile_paras=f"-B /usr/lib/gcc/x86_64-kylin-linux/12/ -fopenacc -fopenmp-targets=amdgcn-amd-amdhsa -Xopenmp-target=amdgcn-amd-amdhsa -march={gfx_arch}")

        # 执行trace测试
        ret, output = exec_cmd(f"hipprof --omp-trace {exe_path}", cwd=test_env)
        assert ret is True
        json_path = self.get_all_prof_files_info(output)["json_list"][0]["name"]

        # 判断不同事件的id不同不同
        event_id_list = []
        for item_d in self.parse_prof_json(json_path):
            if item_d.get("cat") == "DataFlow":
                cur_id = int(item_d.get("id"))
                event_id_list.append(cur_id)
        assert len(event_id_list) > 2, f"Event data was not obtained from the json file({json_path})"
        assert all(count==2 for count in Counter(event_id_list).values())  # 每个事件的id有2个json(1条标题，1条数据)

